package com.example.instrumentation

import org.gradle.api.DefaultTask
import org.gradle.api.file.Directory
import org.gradle.api.file.RegularFile
import org.gradle.api.file.RegularFileProperty
import org.gradle.api.provider.ListProperty
import org.gradle.api.tasks.InputFiles
import org.gradle.api.tasks.OutputFile
import org.gradle.api.tasks.TaskAction
import org.objectweb.asm.ClassReader
import org.objectweb.asm.ClassVisitor
import org.objectweb.asm.ClassWriter
import org.objectweb.asm.MethodVisitor
import org.objectweb.asm.Opcodes
import java.io.File
import java.util.jar.JarEntry
import java.util.jar.JarFile
import java.util.jar.JarOutputStream
import java.util.zip.Deflater


/**
 * create by liaohailong
 * 2025/10/31 23:46
 * desc:
 */
abstract class InmobiTransformTask : DefaultTask() {

    @get:InputFiles
    abstract val inputJars: ListProperty<RegularFile>

    @get:InputFiles
    abstract val inputDirectories: ListProperty<Directory>

    @get:OutputFile
    abstract val transformOutput: RegularFileProperty

    @TaskAction
    fun taskAction() {
        log("taskAction start")
        val outputFile = transformOutput.get().asFile
        if (outputFile.exists()) outputFile.delete() // clear cache
        outputFile.parentFile.mkdirs()
        log("clear cache: ${outputFile.absolutePath}")

        outputFile.outputStream().use { fos ->
            JarOutputStream(fos).use { jos ->
                // this is the fastest mode
                jos.setLevel(Deflater.NO_COMPRESSION)

                // let's go
                transform(jos)
            }
        }

        log("taskAction over")
    }

    private fun transform(jos: JarOutputStream) {
        val jars = inputJars.get() // jars from project, jars from 3rd part aar(local + cloud maven)
        val dirs = inputDirectories.get()

        // handle jars firstly
        log("handle jars start")
        var count = 0
        jars.forEach { rf ->
            val file = rf.asFile
            JarFile(file).use { jf ->
                // a jar file contains a lot of classes
                jf.entries().iterator().forEach { entry ->
                    val entryName = entry.name
                    if (entryName.endsWith(".class")
                        && entryName.endsWith("module-info.class").not()
                    ) {
                        jf.handleClass(entry,jos)
                        count++
//                        log("[jar] found ${entry.name}")
                    }
                }
            }
        }
        log("handle jars over files[$count]")

        // handle directories that all original sources(java+kotlin) here
        log("handle dirs start")
        count = 0
        dirs.forEach { dir ->
            val rootDir = dir.asFile
            val rootUri = rootDir.toURI()

            // traversal all files in all dirs, and put them into output which is a Jar
            rootDir.walk()
                .filter { it.name != "META-INF" }
                .filter { it.isFile }
                .forEach { file ->
                    val entryName = rootUri.relativize(file.toURI()).path
                        .replace(File.separatorChar, '/')
                    val entry = JarEntry(entryName)
                    val bytes = file.readBytes()
                    jos.putNextEntry(entry)
                    jos.write(bytes)
                    jos.closeEntry()

                    count++
//                    log("[dir] found ${entry.name}")
                }
        }
        log("handle dirs over files[$count]")


    }

    private fun JarFile.handleClass(entry: JarEntry, jos: JarOutputStream) {
        val entryName = entry.name
        val target = "com/thinkup/network/inmobi/InmobiTUInitManager.class"
        if (entryName == target) {
            log("found $target")
            val classBytes = this.getInputStream(entry).use { it.readBytes() }

            // use ASM to parse Java byte code
            val classReader = ClassReader(classBytes)
            val classWriter = ClassWriter(0)
            val classVisitor = object : ClassVisitor(Opcodes.ASM9, classWriter) {
                override fun visitMethod(
                    access: Int,
                    name: String?,
                    descriptor: String?,
                    signature: String?,
                    exceptions: Array<out String>?
                ): MethodVisitor? {
                    if (name == "getTopOnInfoExtraMap" && descriptor == "()Ljava/util/Map;") {
                        log("rewrite function: $name")

                        val mv =
                            classWriter.visitMethod(access, name, descriptor, signature, exceptions)

                        mv.visitCode()

                        // === var1 = new HashMap(); ===
                        mv.visitTypeInsn(Opcodes.NEW, "java/util/HashMap")
                        mv.visitInsn(Opcodes.DUP)
                        mv.visitMethodInsn(
                            Opcodes.INVOKESPECIAL,
                            "java/util/HashMap",
                            "<init>",
                            "()V",
                            false
                        )
                        mv.visitVarInsn(Opcodes.ASTORE, 1)

                        // === var1.put("tp-ver", TUSDK.getSDKVersionName()); ===
                        mv.visitVarInsn(Opcodes.ALOAD, 1)
                        mv.visitLdcInsn("tp-ver")
                        mv.visitMethodInsn(
                            Opcodes.INVOKESTATIC,
                            "com/thinkup/core/api/TUSDK",
                            "getSDKVersionName",
                            "()Ljava/lang/String;",
                            false
                        )
                        mv.visitMethodInsn(
                            Opcodes.INVOKEVIRTUAL,
                            "java/util/HashMap",
                            "put",
                            "(Ljava/lang/Object;Ljava/lang/Object;)Ljava/lang/Object;",
                            false
                        )
                        mv.visitInsn(Opcodes.POP)

                        // === return var1; ===
                        mv.visitVarInsn(Opcodes.ALOAD, 1)
                        mv.visitInsn(Opcodes.ARETURN)

                        // === stack and local variable size ==
                        mv.visitMaxs(3, 2)
                        mv.visitEnd()

                        return null
                    }

                    return super.visitMethod(access, name, descriptor, signature, exceptions)
                }
            }
            classReader.accept(classVisitor, 0)

            // get modified Java byte code
            val modifiedBytes = classWriter.toByteArray()
            // new a JarEntry because you maybe modified the length of byte code
            val newEntry = JarEntry(entryName)
            jos.putNextEntry(newEntry)
            jos.write(modifiedBytes)
            jos.closeEntry()
        } else {
            // don't modified just return the original Java byte code
            val bytes = this.getInputStream(entry).use { it.readBytes() }
            jos.putNextEntry(entry)
            jos.write(bytes)
            jos.closeEntry()
        }

    }

    private fun log(msg: String) {
        println("InmobiTransformTask: $msg")
    }

}